"""
计算引擎实现
"""
from typing import Dict, List, Any, Optional
import logging
import pandas as pd
import numpy as np

from .base_engine import BaseEngine
from ..data.loaders import BaseLoader
from ..data.processors import DataCleaner, DataTransformer

logger = logging.getLogger(__name__)


class CalculationEngine(BaseEngine):
    """
    计算引擎，负责执行具体计算任务
    """
    
    def __init__(self, config: Dict[str, Any]):
        """
        初始化计算引擎
        
        Args:
            config: 引擎配置
        """
        super().__init__(config)
        
        # 数据加载器
        self.data_loader = None
        
        # 数据处理器
        self.data_cleaner = DataCleaner()
        self.data_transformer = DataTransformer()
        
        # 计算状态
        self._progress = 0
        self._is_cancelled = False
    
    def validate_input(self, input_data: Dict[str, Any]) -> Dict[str, Any]:
        """
        验证输入数据
        
        Args:
            input_data: 输入数据
            
        Returns:
            验证结果
        """
        errors = []
        
        # 检查必填参数
        required_params = self.get_supported_parameters()
        for param in required_params:
            if param["required"] and param["name"] not in input_data:
                errors.append(f"缺少必填参数: {param['name']}")
        
        # 检查参数类型
        for param_name, param_value in input_data.items():
            param_def = next((p for p in required_params if p["name"] == param_name), None)
            if param_def:
                expected_type = param_def["type"]
                if not isinstance(param_value, eval(expected_type)):
                    errors.append(f"参数 {param_name} 类型错误，应为 {expected_type}")
        
        return {
            "valid": len(errors) == 0,
            "errors": errors
        }
    
    def execute(self, input_data: Dict[str, Any]) -> Dict[str, Any]:
        """
        执行计算
        
        Args:
            input_data: 输入数据
            
        Returns:
            计算结果
        """
        # 重置状态
        self._progress = 0
        self._is_cancelled = False
        
        try:
            # 验证输入
            validation_result = self.validate_input(input_data)
            if not validation_result["valid"]:
                raise ValueError("输入数据验证失败: " + ", ".join(validation_result["errors"]))
            
            # 初始化数据加载器
            self._progress = 5
            self.data_loader = self._init_data_loader(input_data)
            
            # 加载数据
            self._progress = 10
            data = self.data_loader.load()
            
            # 数据清洗
            self._progress = 30
            if "cleaning_rules" in input_data:
                data = self.data_cleaner.clean(data, input_data["cleaning_rules"])
            
            # 数据转换
            self._progress = 50
            if "transformation_rules" in input_data:
                data = self.data_transformer.transform(data, input_data["transformation_rules"])
            
            # 执行计算
            self._progress = 70
            result = self._perform_calculation(data, input_data)
            
            # 完成
            self._progress = 100
            
            return {
                "status": "success",
                "result": result,
                "input_data": input_data,
                "metadata": {
                    "data_shape": data.shape if isinstance(data, pd.DataFrame) else str(data.shape),
                    "execution_time": ""  # 实际实现中应计算执行时间
                }
            }
            
        except Exception as e:
            logger.error(f"计算执行失败: {str(e)}", exc_info=True)
            return {
                "status": "error",
                "error": str(e),
                "input_data": input_data
            }
    
    def get_progress(self) -> float:
        """
        获取计算进度
        
        Returns:
            进度百分比
        """
        return self._progress
    
    def cancel(self) -> bool:
        """
        取消计算
        
        Returns:
            是否取消成功
        """
        self._is_cancelled = True
        return True
    
    def get_supported_parameters(self) -> List[Dict[str, Any]]:
        """
        获取支持的参数列表
        
        Returns:
            参数列表
        """
        return [
            {
                "name": "data_source",
                "type": "dict",
                "description": "数据源配置",
                "required": True
            },
            {
                "name": "cleaning_rules",
                "type": "list",
                "description": "数据清洗规则",
                "required": False
            },
            {
                "name": "transformation_rules",
                "type": "list",
                "description": "数据转换规则",
                "required": False
            },
            {
                "name": "calculation_method",
                "type": "str",
                "description": "计算方法",
                "required": True
            }
        ]
    
    def _init_data_loader(self, input_data: Dict[str, Any]) -> BaseLoader:
        """
        初始化数据加载器
        
        Args:
            input_data: 输入数据
            
        Returns:
            数据加载器实例
        """
        from ..data.loaders import FileLoader, DatabaseLoader, ManualLoader
        
        data_source = input_data["data_source"]
        source_type = data_source.get("type", "file")
        
        if source_type == "file":
            return FileLoader(data_source)
        elif source_type == "database":
            return DatabaseLoader(data_source)
        elif source_type == "manual":
            return ManualLoader(data_source)
        else:
            raise ValueError(f"不支持的数据源类型: {source_type}")
    
    def _perform_calculation(self, data: pd.DataFrame, input_data: Dict[str, Any]) -> Dict[str, Any]:
        """
        执行具体计算
        
        Args:
            data: 输入数据
            input_data: 输入参数
            
        Returns:
            计算结果
        """
        method = input_data["calculation_method"]
        
        if method == "sum":
            # 求和计算
            numeric_cols = data.select_dtypes(include=np.number).columns
            return {
                "type": "sum",
                "results": data[numeric_cols].sum().to_dict()
            }
        
        elif method == "average":
            # 平均值计算
            numeric_cols = data.select_dtypes(include=np.number).columns
            return {
                "type": "average",
                "results": data[numeric_cols].mean().to_dict()
            }
        
        elif method == "correlation":
            # 相关性计算
            numeric_cols = data.select_dtypes(include=np.number).columns
            corr_matrix = data[numeric_cols].corr()
            return {
                "type": "correlation",
                "results": corr_matrix.to_dict()
            }
        
        elif method == "regression":
            # 线性回归
            from sklearn.linear_model import LinearRegression
            
            # 获取自变量和因变量
            x_cols = input_data.get("independent_variables", [])
            y_col = input_data.get("dependent_variable", "")
            
            if not x_cols or not y_col:
                raise ValueError("回归分析需要指定自变量和因变量")
            
            # 准备数据
            X = data[x_cols].values
            y = data[y_col].values
            
            # 训练模型
            model = LinearRegression()
            model.fit(X, y)
            
            # 返回结果
            return {
                "type": "regression",
                "coefficients": dict(zip(x_cols, model.coef_)),
                "intercept": model.intercept_,
                "r_squared": model.score(X, y)
            }
        
        else:
            raise ValueError(f"不支持的计算方法: {method}")